import requests
import concurrent.futures
from datetime import datetime


class PrometheusClient:
    def __init__(self, base_url: str, timeout: int = 15):  # 增加超时时间
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.timeout = timeout

    def query(self, query: str):
        try:
            resp = self.session.get(
                f"{self.base_url}/api/v1/query",
                params={"query": query},
                timeout=self.timeout
            )
            resp.raise_for_status()
            data = resp.json()
            return data.get("data", {}).get("result", []) if data.get("status") == "success" else []
        except Exception as e:
            print(f"[查询失败] {e} | query={query}")
            return []

    def get_deployment_creation_time(self, deployment: str, namespace: str):
        """获取deployment的创建时间"""
        query = f'kube_pod_created{{pod=~"{deployment}-[a-z0-9]+-[a-z0-9]+", namespace="{namespace}"}}'
        result = self.query(query)

        if result:
            # 返回第一个结果的时间戳（Unix时间戳）
            return float(result[0].get('value', [0, 0])[1])
        return None

    def get_images_from_env(self, deployment: str, namespace: str):
        """从指定环境和命名空间获取deployment的镜像"""
        # 使用精确匹配查询，避免模糊匹配到其他deployment
        images = set()

        # 查询deployment镜像
        deployment_query = f'kube_deployment_container_image{{namespace="{namespace}", deployment="{deployment}"}}'
        deployment_result = self.query(deployment_query)

        for item in deployment_result:
            image = item.get("metric", {}).get("image")
            if image:
                images.add(image)

        # 如果没有找到，尝试通过pod查询（也使用精确匹配）
        if not images:
            pod_query = f'kube_pod_container_info{{namespace="{namespace}", pod=~"{deployment}-[a-z0-9]+-[a-z0-9]+"}}'
            pod_result = self.query(pod_query)

            for item in pod_result:
                image = item.get("metric", {}).get("image")
                if image:
                    images.add(image)

        # 清理镜像名称
        cleaned_images = [self.clean_image_name(img) for img in images] if images else ["Not Found"]
        return sorted(cleaned_images)

    def clean_image_name(self, image: str) -> str:
        """清理镜像名称，去除指定的前缀"""
        prefixes_to_remove = [
            "ksher-ccr.tencentcloudcr.com/ksher-prod/",
            "test-ksher-ccr.tencentcloudcr.com/ksher-uat/",
            "test-ksher-ccr.tencentcloudcr.com/ksher-test/",
            "test-ksher-ccr.tencentcloudcr.com/ksher-dev/",
        ]

        for prefix in prefixes_to_remove:
            if image.startswith(prefix):
                return image[len(prefix):]

        return image

    def query_batch(self, queries: list):
        """批量执行查询"""
        results = {}
        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            future_to_query = {executor.submit(self.query, q): q for q in queries}
            for future in concurrent.futures.as_completed(future_to_query):
                query = future_to_query[future]
                try:
                    results[query] = future.result()
                except Exception as e:
                    print(f"Batch query failed for {query}: {e}")
                    results[query] = []
        return results